{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from fastai.vision import *\n",
    "import torch\n",
    "import warnings\n",
    "from torchsummary import summary\n",
    "# from models.custom_resnet import _resnet, Bottleneck\n",
    "from models.resnet_cifar import *\n",
    "torch.cuda.set_device(0)\n",
    "torch.manual_seed(1)\n",
    "torch.cuda.manual_seed(1)\n",
    "\n",
    "# stage should be in 0 to 5 (5 for classifier stage)\n",
    "hyper_params = {\n",
    "    \"stage\": 5,\n",
    "    \"repeated\": 0,\n",
    "    \"num_classes\": 10,\n",
    "    \"batch_size\": 16,\n",
    "    \"num_epochs\": 2,\n",
    "    \"learning_rate\": 1e-4\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = untar_data(URLs.IMAGENETTE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tfms = get_transforms(do_flip=False)\n",
    "data = ImageDataBunch.from_folder(path, train = 'train', valid = 'val', bs = hyper_params[\"batch_size\"], size = 224, ds_tfms = tfms).normalize(imagenet_stats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model on GPU\n"
     ]
    }
   ],
   "source": [
    "learn = cnn_learner(data, models.resnet50, metrics = accuracy, pretrained = True)\n",
    "learn.load('resnet50_imagenette_bs32')\n",
    "learn.freeze()\n",
    "# learn.summary()\n",
    "\n",
    "net = _resnet('resnet50', Bottleneck, [2, 2, 2, 1], pretrained = False, progress = False)\n",
    "\n",
    "if torch.cuda.is_available() : \n",
    "    net = net.cuda()\n",
    "    print('Model on GPU')\n",
    "    \n",
    "# for name, param in net.named_parameters() : \n",
    "#     print(name, param.shape)\n",
    "#     param.requires_grad = False\n",
    "#     if name[5] == str(hyper_params['stage']) and hyper_params['stage'] != 0 :\n",
    "#         param.requires_grad = True\n",
    "#     elif (name[0] == 'b' or name[0] == 'c') and hyper_params['stage'] == 0 : \n",
    "#         param.requires_grad = True\n",
    "#     elif name[0] == 'f' and hyper_params['stage'] == 5 : \n",
    "#         param.requires_grad = True\n",
    "#     print(param.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace)\n",
       "      (downsample): Sequential(\n",
       "        (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveConcatPool2d(\n",
       "    (ap): AdaptiveAvgPool2d(output_size=1)\n",
       "    (mp): AdaptiveMaxPool2d(output_size=1)\n",
       "  )\n",
       "  (fc): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "  (relu2): ReLU(inplace)\n",
       "  (fc2): Linear(in_features=1024, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 112, 112]             128\n",
      "              ReLU-3         [-1, 64, 112, 112]               0\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5           [-1, 64, 56, 56]           4,096\n",
      "       BatchNorm2d-6           [-1, 64, 56, 56]             128\n",
      "              ReLU-7           [-1, 64, 56, 56]               0\n",
      "            Conv2d-8           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-9           [-1, 64, 56, 56]             128\n",
      "             ReLU-10           [-1, 64, 56, 56]               0\n",
      "           Conv2d-11          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-12          [-1, 256, 56, 56]             512\n",
      "           Conv2d-13          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-14          [-1, 256, 56, 56]             512\n",
      "             ReLU-15          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-16          [-1, 256, 56, 56]               0\n",
      "           Conv2d-17           [-1, 64, 56, 56]          16,384\n",
      "      BatchNorm2d-18           [-1, 64, 56, 56]             128\n",
      "             ReLU-19           [-1, 64, 56, 56]               0\n",
      "           Conv2d-20           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-21           [-1, 64, 56, 56]             128\n",
      "             ReLU-22           [-1, 64, 56, 56]               0\n",
      "           Conv2d-23          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-24          [-1, 256, 56, 56]             512\n",
      "             ReLU-25          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-26          [-1, 256, 56, 56]               0\n",
      "           Conv2d-27          [-1, 128, 56, 56]          32,768\n",
      "      BatchNorm2d-28          [-1, 128, 56, 56]             256\n",
      "             ReLU-29          [-1, 128, 56, 56]               0\n",
      "           Conv2d-30          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-31          [-1, 128, 28, 28]             256\n",
      "             ReLU-32          [-1, 128, 28, 28]               0\n",
      "           Conv2d-33          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-34          [-1, 512, 28, 28]           1,024\n",
      "           Conv2d-35          [-1, 512, 28, 28]         131,072\n",
      "      BatchNorm2d-36          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-37          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-38          [-1, 512, 28, 28]               0\n",
      "           Conv2d-39          [-1, 128, 28, 28]          65,536\n",
      "      BatchNorm2d-40          [-1, 128, 28, 28]             256\n",
      "             ReLU-41          [-1, 128, 28, 28]               0\n",
      "           Conv2d-42          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-43          [-1, 128, 28, 28]             256\n",
      "             ReLU-44          [-1, 128, 28, 28]               0\n",
      "           Conv2d-45          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-46          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-47          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-48          [-1, 512, 28, 28]               0\n",
      "           Conv2d-49          [-1, 256, 28, 28]         131,072\n",
      "      BatchNorm2d-50          [-1, 256, 28, 28]             512\n",
      "             ReLU-51          [-1, 256, 28, 28]               0\n",
      "           Conv2d-52          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-53          [-1, 256, 14, 14]             512\n",
      "             ReLU-54          [-1, 256, 14, 14]               0\n",
      "           Conv2d-55         [-1, 1024, 14, 14]         262,144\n",
      "      BatchNorm2d-56         [-1, 1024, 14, 14]           2,048\n",
      "           Conv2d-57         [-1, 1024, 14, 14]         524,288\n",
      "      BatchNorm2d-58         [-1, 1024, 14, 14]           2,048\n",
      "             ReLU-59         [-1, 1024, 14, 14]               0\n",
      "       Bottleneck-60         [-1, 1024, 14, 14]               0\n",
      "           Conv2d-61          [-1, 256, 14, 14]         262,144\n",
      "      BatchNorm2d-62          [-1, 256, 14, 14]             512\n",
      "             ReLU-63          [-1, 256, 14, 14]               0\n",
      "           Conv2d-64          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-65          [-1, 256, 14, 14]             512\n",
      "             ReLU-66          [-1, 256, 14, 14]               0\n",
      "           Conv2d-67         [-1, 1024, 14, 14]         262,144\n",
      "      BatchNorm2d-68         [-1, 1024, 14, 14]           2,048\n",
      "             ReLU-69         [-1, 1024, 14, 14]               0\n",
      "       Bottleneck-70         [-1, 1024, 14, 14]               0\n",
      "           Conv2d-71          [-1, 512, 14, 14]         524,288\n",
      "      BatchNorm2d-72          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-73          [-1, 512, 14, 14]               0\n",
      "           Conv2d-74            [-1, 512, 7, 7]       2,359,296\n",
      "      BatchNorm2d-75            [-1, 512, 7, 7]           1,024\n",
      "             ReLU-76            [-1, 512, 7, 7]               0\n",
      "           Conv2d-77           [-1, 2048, 7, 7]       1,048,576\n",
      "      BatchNorm2d-78           [-1, 2048, 7, 7]           4,096\n",
      "           Conv2d-79           [-1, 2048, 7, 7]       2,097,152\n",
      "      BatchNorm2d-80           [-1, 2048, 7, 7]           4,096\n",
      "             ReLU-81           [-1, 2048, 7, 7]               0\n",
      "       Bottleneck-82           [-1, 2048, 7, 7]               0\n",
      "AdaptiveMaxPool2d-83           [-1, 2048, 1, 1]               0\n",
      "AdaptiveAvgPool2d-84           [-1, 2048, 1, 1]               0\n",
      "AdaptiveConcatPool2d-85           [-1, 4096, 1, 1]               0\n",
      "           Linear-86                 [-1, 1024]       4,195,328\n",
      "             ReLU-87                 [-1, 1024]               0\n",
      "           Linear-88                   [-1, 10]          10,250\n",
      "================================================================\n",
      "Total params: 13,689,162\n",
      "Trainable params: 13,689,162\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 177.13\n",
      "Params size (MB): 52.22\n",
      "Estimated Total Size (MB): 229.92\n",
      "----------------------------------------------------------------\n",
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 112, 112]           9,408\n",
      "       BatchNorm2d-2         [-1, 64, 112, 112]             128\n",
      "              ReLU-3         [-1, 64, 112, 112]               0\n",
      "         MaxPool2d-4           [-1, 64, 56, 56]               0\n",
      "            Conv2d-5           [-1, 64, 56, 56]           4,096\n",
      "       BatchNorm2d-6           [-1, 64, 56, 56]             128\n",
      "              ReLU-7           [-1, 64, 56, 56]               0\n",
      "            Conv2d-8           [-1, 64, 56, 56]          36,864\n",
      "       BatchNorm2d-9           [-1, 64, 56, 56]             128\n",
      "             ReLU-10           [-1, 64, 56, 56]               0\n",
      "           Conv2d-11          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-12          [-1, 256, 56, 56]             512\n",
      "           Conv2d-13          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-14          [-1, 256, 56, 56]             512\n",
      "             ReLU-15          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-16          [-1, 256, 56, 56]               0\n",
      "           Conv2d-17           [-1, 64, 56, 56]          16,384\n",
      "      BatchNorm2d-18           [-1, 64, 56, 56]             128\n",
      "             ReLU-19           [-1, 64, 56, 56]               0\n",
      "           Conv2d-20           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-21           [-1, 64, 56, 56]             128\n",
      "             ReLU-22           [-1, 64, 56, 56]               0\n",
      "           Conv2d-23          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-24          [-1, 256, 56, 56]             512\n",
      "             ReLU-25          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-26          [-1, 256, 56, 56]               0\n",
      "           Conv2d-27           [-1, 64, 56, 56]          16,384\n",
      "      BatchNorm2d-28           [-1, 64, 56, 56]             128\n",
      "             ReLU-29           [-1, 64, 56, 56]               0\n",
      "           Conv2d-30           [-1, 64, 56, 56]          36,864\n",
      "      BatchNorm2d-31           [-1, 64, 56, 56]             128\n",
      "             ReLU-32           [-1, 64, 56, 56]               0\n",
      "           Conv2d-33          [-1, 256, 56, 56]          16,384\n",
      "      BatchNorm2d-34          [-1, 256, 56, 56]             512\n",
      "             ReLU-35          [-1, 256, 56, 56]               0\n",
      "       Bottleneck-36          [-1, 256, 56, 56]               0\n",
      "           Conv2d-37          [-1, 128, 56, 56]          32,768\n",
      "      BatchNorm2d-38          [-1, 128, 56, 56]             256\n",
      "             ReLU-39          [-1, 128, 56, 56]               0\n",
      "           Conv2d-40          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-41          [-1, 128, 28, 28]             256\n",
      "             ReLU-42          [-1, 128, 28, 28]               0\n",
      "           Conv2d-43          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-44          [-1, 512, 28, 28]           1,024\n",
      "           Conv2d-45          [-1, 512, 28, 28]         131,072\n",
      "      BatchNorm2d-46          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-47          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-48          [-1, 512, 28, 28]               0\n",
      "           Conv2d-49          [-1, 128, 28, 28]          65,536\n",
      "      BatchNorm2d-50          [-1, 128, 28, 28]             256\n",
      "             ReLU-51          [-1, 128, 28, 28]               0\n",
      "           Conv2d-52          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-53          [-1, 128, 28, 28]             256\n",
      "             ReLU-54          [-1, 128, 28, 28]               0\n",
      "           Conv2d-55          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-56          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-57          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-58          [-1, 512, 28, 28]               0\n",
      "           Conv2d-59          [-1, 128, 28, 28]          65,536\n",
      "      BatchNorm2d-60          [-1, 128, 28, 28]             256\n",
      "             ReLU-61          [-1, 128, 28, 28]               0\n",
      "           Conv2d-62          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-63          [-1, 128, 28, 28]             256\n",
      "             ReLU-64          [-1, 128, 28, 28]               0\n",
      "           Conv2d-65          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-66          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-67          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-68          [-1, 512, 28, 28]               0\n",
      "           Conv2d-69          [-1, 128, 28, 28]          65,536\n",
      "      BatchNorm2d-70          [-1, 128, 28, 28]             256\n",
      "             ReLU-71          [-1, 128, 28, 28]               0\n",
      "           Conv2d-72          [-1, 128, 28, 28]         147,456\n",
      "      BatchNorm2d-73          [-1, 128, 28, 28]             256\n",
      "             ReLU-74          [-1, 128, 28, 28]               0\n",
      "           Conv2d-75          [-1, 512, 28, 28]          65,536\n",
      "      BatchNorm2d-76          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-77          [-1, 512, 28, 28]               0\n",
      "       Bottleneck-78          [-1, 512, 28, 28]               0\n",
      "           Conv2d-79          [-1, 256, 28, 28]         131,072\n",
      "      BatchNorm2d-80          [-1, 256, 28, 28]             512\n",
      "             ReLU-81          [-1, 256, 28, 28]               0\n",
      "           Conv2d-82          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-83          [-1, 256, 14, 14]             512\n",
      "             ReLU-84          [-1, 256, 14, 14]               0\n",
      "           Conv2d-85         [-1, 1024, 14, 14]         262,144\n",
      "      BatchNorm2d-86         [-1, 1024, 14, 14]           2,048\n",
      "           Conv2d-87         [-1, 1024, 14, 14]         524,288\n",
      "      BatchNorm2d-88         [-1, 1024, 14, 14]           2,048\n",
      "             ReLU-89         [-1, 1024, 14, 14]               0\n",
      "       Bottleneck-90         [-1, 1024, 14, 14]               0\n",
      "           Conv2d-91          [-1, 256, 14, 14]         262,144\n",
      "      BatchNorm2d-92          [-1, 256, 14, 14]             512\n",
      "             ReLU-93          [-1, 256, 14, 14]               0\n",
      "           Conv2d-94          [-1, 256, 14, 14]         589,824\n",
      "      BatchNorm2d-95          [-1, 256, 14, 14]             512\n",
      "             ReLU-96          [-1, 256, 14, 14]               0\n",
      "           Conv2d-97         [-1, 1024, 14, 14]         262,144\n",
      "      BatchNorm2d-98         [-1, 1024, 14, 14]           2,048\n",
      "             ReLU-99         [-1, 1024, 14, 14]               0\n",
      "      Bottleneck-100         [-1, 1024, 14, 14]               0\n",
      "          Conv2d-101          [-1, 256, 14, 14]         262,144\n",
      "     BatchNorm2d-102          [-1, 256, 14, 14]             512\n",
      "            ReLU-103          [-1, 256, 14, 14]               0\n",
      "          Conv2d-104          [-1, 256, 14, 14]         589,824\n",
      "     BatchNorm2d-105          [-1, 256, 14, 14]             512\n",
      "            ReLU-106          [-1, 256, 14, 14]               0\n",
      "          Conv2d-107         [-1, 1024, 14, 14]         262,144\n",
      "     BatchNorm2d-108         [-1, 1024, 14, 14]           2,048\n",
      "            ReLU-109         [-1, 1024, 14, 14]               0\n",
      "      Bottleneck-110         [-1, 1024, 14, 14]               0\n",
      "          Conv2d-111          [-1, 256, 14, 14]         262,144\n",
      "     BatchNorm2d-112          [-1, 256, 14, 14]             512\n",
      "            ReLU-113          [-1, 256, 14, 14]               0\n",
      "          Conv2d-114          [-1, 256, 14, 14]         589,824\n",
      "     BatchNorm2d-115          [-1, 256, 14, 14]             512\n",
      "            ReLU-116          [-1, 256, 14, 14]               0\n",
      "          Conv2d-117         [-1, 1024, 14, 14]         262,144\n",
      "     BatchNorm2d-118         [-1, 1024, 14, 14]           2,048\n",
      "            ReLU-119         [-1, 1024, 14, 14]               0\n",
      "      Bottleneck-120         [-1, 1024, 14, 14]               0\n",
      "          Conv2d-121          [-1, 256, 14, 14]         262,144\n",
      "     BatchNorm2d-122          [-1, 256, 14, 14]             512\n",
      "            ReLU-123          [-1, 256, 14, 14]               0\n",
      "          Conv2d-124          [-1, 256, 14, 14]         589,824\n",
      "     BatchNorm2d-125          [-1, 256, 14, 14]             512\n",
      "            ReLU-126          [-1, 256, 14, 14]               0\n",
      "          Conv2d-127         [-1, 1024, 14, 14]         262,144\n",
      "     BatchNorm2d-128         [-1, 1024, 14, 14]           2,048\n",
      "            ReLU-129         [-1, 1024, 14, 14]               0\n",
      "      Bottleneck-130         [-1, 1024, 14, 14]               0\n",
      "          Conv2d-131          [-1, 256, 14, 14]         262,144\n",
      "     BatchNorm2d-132          [-1, 256, 14, 14]             512\n",
      "            ReLU-133          [-1, 256, 14, 14]               0\n",
      "          Conv2d-134          [-1, 256, 14, 14]         589,824\n",
      "     BatchNorm2d-135          [-1, 256, 14, 14]             512\n",
      "            ReLU-136          [-1, 256, 14, 14]               0\n",
      "          Conv2d-137         [-1, 1024, 14, 14]         262,144\n",
      "     BatchNorm2d-138         [-1, 1024, 14, 14]           2,048\n",
      "            ReLU-139         [-1, 1024, 14, 14]               0\n",
      "      Bottleneck-140         [-1, 1024, 14, 14]               0\n",
      "          Conv2d-141          [-1, 512, 14, 14]         524,288\n",
      "     BatchNorm2d-142          [-1, 512, 14, 14]           1,024\n",
      "            ReLU-143          [-1, 512, 14, 14]               0\n",
      "          Conv2d-144            [-1, 512, 7, 7]       2,359,296\n",
      "     BatchNorm2d-145            [-1, 512, 7, 7]           1,024\n",
      "            ReLU-146            [-1, 512, 7, 7]               0\n",
      "          Conv2d-147           [-1, 2048, 7, 7]       1,048,576\n",
      "     BatchNorm2d-148           [-1, 2048, 7, 7]           4,096\n",
      "          Conv2d-149           [-1, 2048, 7, 7]       2,097,152\n",
      "     BatchNorm2d-150           [-1, 2048, 7, 7]           4,096\n",
      "            ReLU-151           [-1, 2048, 7, 7]               0\n",
      "      Bottleneck-152           [-1, 2048, 7, 7]               0\n",
      "          Conv2d-153            [-1, 512, 7, 7]       1,048,576\n",
      "     BatchNorm2d-154            [-1, 512, 7, 7]           1,024\n",
      "            ReLU-155            [-1, 512, 7, 7]               0\n",
      "          Conv2d-156            [-1, 512, 7, 7]       2,359,296\n",
      "     BatchNorm2d-157            [-1, 512, 7, 7]           1,024\n",
      "            ReLU-158            [-1, 512, 7, 7]               0\n",
      "          Conv2d-159           [-1, 2048, 7, 7]       1,048,576\n",
      "     BatchNorm2d-160           [-1, 2048, 7, 7]           4,096\n",
      "            ReLU-161           [-1, 2048, 7, 7]               0\n",
      "      Bottleneck-162           [-1, 2048, 7, 7]               0\n",
      "          Conv2d-163            [-1, 512, 7, 7]       1,048,576\n",
      "     BatchNorm2d-164            [-1, 512, 7, 7]           1,024\n",
      "            ReLU-165            [-1, 512, 7, 7]               0\n",
      "          Conv2d-166            [-1, 512, 7, 7]       2,359,296\n",
      "     BatchNorm2d-167            [-1, 512, 7, 7]           1,024\n",
      "            ReLU-168            [-1, 512, 7, 7]               0\n",
      "          Conv2d-169           [-1, 2048, 7, 7]       1,048,576\n",
      "     BatchNorm2d-170           [-1, 2048, 7, 7]           4,096\n",
      "            ReLU-171           [-1, 2048, 7, 7]               0\n",
      "      Bottleneck-172           [-1, 2048, 7, 7]               0\n",
      "AdaptiveMaxPool2d-173           [-1, 2048, 1, 1]               0\n",
      "AdaptiveAvgPool2d-174           [-1, 2048, 1, 1]               0\n",
      "AdaptiveConcatPool2d-175           [-1, 4096, 1, 1]               0\n",
      "         Flatten-176                 [-1, 4096]               0\n",
      "     BatchNorm1d-177                 [-1, 4096]           8,192\n",
      "         Dropout-178                 [-1, 4096]               0\n",
      "          Linear-179                  [-1, 512]       2,097,664\n",
      "            ReLU-180                  [-1, 512]               0\n",
      "     BatchNorm1d-181                  [-1, 512]           1,024\n",
      "         Dropout-182                  [-1, 512]               0\n",
      "          Linear-183                   [-1, 10]           5,130\n",
      "================================================================\n",
      "Total params: 25,620,042\n",
      "Trainable params: 2,165,130\n",
      "Non-trainable params: 23,454,912\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 286.71\n",
      "Params size (MB): 97.73\n",
      "Estimated Total Size (MB): 385.01\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# x, y = next(iter(data.train_dl))\n",
    "# net(torch.autograd.Variable(x).cuda())\n",
    "summary(net, (3, 224, 224))\n",
    "summary(learn.model, (3, 224, 224))\n",
    "# print(learn.summary())\n",
    "# net.layer1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SaveFeatures :\n",
    "    def __init__(self, m) : \n",
    "        self.handle = m.register_forward_hook(self.hook_fn)\n",
    "    def hook_fn(self, m, inp, outp) : \n",
    "        self.features = outp\n",
    "    def remove(self) :\n",
    "        self.handle.remove()\n",
    "        \n",
    "mdl = learn.model\n",
    "del learn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "sf = [SaveFeatures(m) for m in [mdl[0][2], mdl[0][4], mdl[0][5], mdl[0][6], mdl[0][7]]]\n",
    "sf2 = [SaveFeatures(m) for m in [net.relu2, net.layer1, net.layer2, net.layer3, net.layer4]]\n",
    "\n",
    "### For testing hooks \n",
    "# x, y = next(iter(data.train_dl))\n",
    "# x = torch.autograd.Variable(x).cuda()\n",
    "# out1 = mdl(x)\n",
    "# out2 = net(x)\n",
    "# for i in range(4) : \n",
    "#     print('resnet50 : ', sf[i].features.shape)\n",
    "#     print('custom model : ', sf2[i].features.shape)\n",
    "# del x, y, sf, sf2\n",
    "# del out1, out2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Stage-wise training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch =  1  step =  50  of total steps  805  loss =  0.15482433140277863\n",
      "epoch =  1  step =  100  of total steps  805  loss =  0.12196048349142075\n",
      "epoch =  1  step =  150  of total steps  805  loss =  0.14741818606853485\n",
      "epoch =  1  step =  200  of total steps  805  loss =  0.156645268201828\n",
      "epoch =  1  step =  250  of total steps  805  loss =  0.1329699158668518\n",
      "epoch =  1  step =  300  of total steps  805  loss =  0.1337122768163681\n",
      "epoch =  1  step =  350  of total steps  805  loss =  0.11640556156635284\n",
      "epoch =  1  step =  400  of total steps  805  loss =  0.11518991738557816\n",
      "epoch =  1  step =  450  of total steps  805  loss =  0.11609910428524017\n",
      "epoch =  1  step =  500  of total steps  805  loss =  0.11605231463909149\n",
      "epoch =  1  step =  550  of total steps  805  loss =  0.10304179042577744\n",
      "epoch =  1  step =  600  of total steps  805  loss =  0.10243804007768631\n",
      "epoch =  1  step =  650  of total steps  805  loss =  0.10527341067790985\n",
      "epoch =  1  step =  700  of total steps  805  loss =  0.10937318950891495\n",
      "epoch =  1  step =  750  of total steps  805  loss =  0.08963960409164429\n",
      "epoch =  1  step =  800  of total steps  805  loss =  0.09278739243745804\n",
      "epoch :  1  /  2  | TL :  0.12025123917908402  | VL :  0.08692677388899028\n",
      "saving model\n",
      "epoch =  2  step =  50  of total steps  805  loss =  0.08528214693069458\n",
      "epoch =  2  step =  100  of total steps  805  loss =  0.08195324242115021\n",
      "epoch =  2  step =  150  of total steps  805  loss =  0.07963912934064865\n",
      "epoch =  2  step =  200  of total steps  805  loss =  0.08056091517210007\n",
      "epoch =  2  step =  250  of total steps  805  loss =  0.08498793095350266\n",
      "epoch =  2  step =  300  of total steps  805  loss =  0.08090215921401978\n",
      "epoch =  2  step =  350  of total steps  805  loss =  0.07447414845228195\n",
      "epoch =  2  step =  400  of total steps  805  loss =  0.06504533439874649\n",
      "epoch =  2  step =  450  of total steps  805  loss =  0.06957031041383743\n",
      "epoch =  2  step =  500  of total steps  805  loss =  0.06390005350112915\n",
      "epoch =  2  step =  550  of total steps  805  loss =  0.06317362189292908\n",
      "epoch =  2  step =  600  of total steps  805  loss =  0.06319869309663773\n",
      "epoch =  2  step =  650  of total steps  805  loss =  0.06257843226194382\n",
      "epoch =  2  step =  700  of total steps  805  loss =  0.05500667914748192\n",
      "epoch =  2  step =  750  of total steps  805  loss =  0.06255189329385757\n",
      "epoch =  2  step =  800  of total steps  805  loss =  0.05196793004870415\n",
      "epoch :  2  /  2  | TL :  0.07195023581940935  | VL :  0.055773036438040435\n",
      "saving model\n"
     ]
    }
   ],
   "source": [
    "if hyper_params['stage'] == 0 : \n",
    "    filename = '../saved_models/large_stage' + str(hyper_params['stage']) + '/model' + str(hyper_params['repeated']) + '.pt'\n",
    "else : \n",
    "    filename = '../saved_models/large_stage' + str(hyper_params['stage'] + 1) + '/model' + str(hyper_params['repeated']) + '.pt'\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr = hyper_params[\"learning_rate\"])\n",
    "total_step = len(data.train_ds) // hyper_params[\"batch_size\"]\n",
    "train_loss_list = list()\n",
    "val_loss_list = list()\n",
    "min_val = 100\n",
    "for epoch in range(hyper_params[\"num_epochs\"]):\n",
    "    trn = []\n",
    "    net.train()\n",
    "    for i, (images, labels) in enumerate(data.train_dl) :\n",
    "        if torch.cuda.is_available():\n",
    "            images = torch.autograd.Variable(images).cuda().float()\n",
    "            labels = torch.autograd.Variable(labels).cuda()\n",
    "        else : \n",
    "            images = torch.autograd.Variable(images).float()\n",
    "            labels = torch.autograd.Variable(labels)\n",
    "\n",
    "        y_pred = net(images)\n",
    "        y_pred2 = mdl(images)\n",
    "\n",
    "        loss = F.mse_loss(sf2[hyper_params[\"stage\"]].features, sf[hyper_params[\"stage\"]].features)\n",
    "        trn.append(loss.item())\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "#         torch.nn.utils.clip_grad_value_(net.parameters(), 10)\n",
    "        optimizer.step()\n",
    "\n",
    "        if i % 50 == 49 :\n",
    "            print('epoch = ', epoch + 1, ' step = ', i + 1, ' of total steps ', total_step, ' loss = ', loss.item())\n",
    "\n",
    "    train_loss = (sum(trn) / len(trn))\n",
    "    train_loss_list.append(train_loss)\n",
    "\n",
    "    net.eval()\n",
    "    val = []\n",
    "    with torch.no_grad() :\n",
    "        for i, (images, labels) in enumerate(data.valid_dl) :\n",
    "            if torch.cuda.is_available():\n",
    "                images = torch.autograd.Variable(images).cuda().float()\n",
    "                labels = torch.autograd.Variable(labels).cuda()\n",
    "            else : \n",
    "                images = torch.autograd.Variable(images).float()\n",
    "                labels = torch.autograd.Variable(labels)\n",
    "\n",
    "            # Forward pass\n",
    "            y_pred = net(images)\n",
    "            y_pred2 = mdl(images)\n",
    "            loss = F.mse_loss(sf[hyper_params[\"stage\"]].features, sf2[hyper_params[\"stage\"]].features)\n",
    "            val.append(loss.item())\n",
    "\n",
    "    val_loss = sum(val) / len(val)\n",
    "    val_loss_list.append(val_loss)\n",
    "    print('epoch : ', epoch + 1, ' / ', hyper_params[\"num_epochs\"], ' | TL : ', train_loss, ' | VL : ', val_loss)\n",
    "\n",
    "    if val_loss < min_val :\n",
    "        print('saving model')\n",
    "        min_val = val_loss\n",
    "        torch.save(net.state_dict(), filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# learn = Learner(data, net, metrics = accuracy)\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/stage5/model0.pt', map_location = 'cpu'))\n",
    "net = net.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(range(hyper_params[\"num_epochs\"]), train_loss_list, label = 'train_loss') \n",
    "plt.plot(range(hyper_params[\"num_epochs\"]), val_loss_list, label = 'val_loss')\n",
    "plt.legend()\n",
    "plt.savefig('../figures/stage5/train_loss.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training the classifier only after stage-wise training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hyper_params = {\n",
    "    \"stage\": 5,\n",
    "    \"repeated\": 3,\n",
    "    \"num_classes\": 10,\n",
    "    \"batch_size\": 64,\n",
    "    \"num_epochs\": 100,\n",
    "    \"learning_rate\": 1e-4\n",
    "}\n",
    "\n",
    "class Flatten(nn.Module) :\n",
    "    def forward(self, input):\n",
    "        return input.view(input.size(0), -1)\n",
    "\n",
    "def conv2(ni, nf) : \n",
    "    return conv_layer(ni, nf, stride = 2)\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(self, nf):\n",
    "        super().__init__()\n",
    "        self.conv1 = conv_layer(nf,nf)\n",
    "        \n",
    "    def forward(self, x): \n",
    "        return (x + self.conv1(x))\n",
    "\n",
    "def conv_and_res(ni, nf): \n",
    "    return nn.Sequential(conv2(ni, nf), ResBlock(nf))\n",
    "\n",
    "def conv_(nf) : \n",
    "    return nn.Sequential(conv_layer(nf, nf), ResBlock(nf))\n",
    "    \n",
    "net = nn.Sequential(\n",
    "    conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "    nn.MaxPool2d(3, 2, padding = 1),\n",
    "    conv_(64),\n",
    "    conv_and_res(64, 128),\n",
    "    conv_and_res(128, 256),\n",
    "    AdaptiveConcatPool2d(),\n",
    "    Flatten(),\n",
    "    nn.Linear(2 * 256, 128),\n",
    "    nn.Linear(128, hyper_params[\"num_classes\"])\n",
    ")\n",
    "\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/small_stage4/model1.pt', map_location = 'cpu'))\n",
    "\n",
    "if torch.cuda.is_available() : \n",
    "    net = net.cuda()\n",
    "    print('Model on GPU')\n",
    "    \n",
    "for name, param in net.named_parameters() : \n",
    "    print(name, param.shape)\n",
    "    param.requires_grad = False\n",
    "    if name[0] == '7' or name[0] == '8':\n",
    "        param.requires_grad = True\n",
    "    print(param.requires_grad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.982\n",
      "0.934\n"
     ]
    }
   ],
   "source": [
    "def _get_accuracy(dataloader, Net):\n",
    "    total = 0\n",
    "    correct = 0\n",
    "    Net.eval()\n",
    "    for i, (images, labels) in enumerate(dataloader):\n",
    "        images = torch.autograd.Variable(images).float()\n",
    "        labels = torch.autograd.Variable(labels).float()\n",
    "        \n",
    "        if torch.cuda.is_available() : \n",
    "            images = images.cuda()\n",
    "            labels = labels.cuda()\n",
    "\n",
    "        outputs = Net.forward(images)\n",
    "        outputs = F.log_softmax(outputs, dim = 1)\n",
    "\n",
    "        _, pred_ind = torch.max(outputs, 1)\n",
    "        \n",
    "        # converting to numpy arrays\n",
    "        labels = labels.data.cpu().numpy()\n",
    "        pred_ind = pred_ind.data.cpu().numpy()\n",
    "        \n",
    "        # get difference\n",
    "        diff_ind = labels - pred_ind\n",
    "        # correctly classified will be 1 and will get added\n",
    "        # incorrectly classified will be 0\n",
    "        correct += np.count_nonzero(diff_ind == 0)\n",
    "        total += len(diff_ind)\n",
    "\n",
    "    accuracy = correct / total\n",
    "    # print(len(diff_ind))\n",
    "    return accuracy\n",
    "\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/large_classifier/model0.pt', map_location = 'cpu'))\n",
    "net.cuda()\n",
    "print(_get_accuracy(data.valid_dl, net))\n",
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/large_no_teacher/model0.pt', map_location = 'cpu'))\n",
    "net.cuda()\n",
    "print(_get_accuracy(data.valid_dl, net))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(net.parameters(), lr = hyper_params[\"learning_rate\"])\n",
    "total_step = len(data.train_ds) // hyper_params[\"batch_size\"]\n",
    "train_loss_list = list()\n",
    "val_loss_list = list()\n",
    "min_val = 0\n",
    "for epoch in range(hyper_params[\"num_epochs\"]):\n",
    "    trn = []\n",
    "    net.train()\n",
    "    for i, (images, labels) in enumerate(data.train_dl) :\n",
    "        if torch.cuda.is_available():\n",
    "            images = torch.autograd.Variable(images).cuda().float()\n",
    "            labels = torch.autograd.Variable(labels).cuda()\n",
    "        else : \n",
    "            images = torch.autograd.Variable(images).float()\n",
    "            labels = torch.autograd.Variable(labels)\n",
    "\n",
    "        y_pred = net(images)\n",
    "\n",
    "        loss = F.cross_entropy(y_pred, labels)\n",
    "        trn.append(loss.item())\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "#         torch.nn.utils.clip_grad_value_(net.parameters(), 10)\n",
    "        optimizer.step()\n",
    "\n",
    "        if i % 50 == 49 :\n",
    "            print('epoch = ', epoch, ' step = ', i + 1, ' of total steps ', total_step, ' loss = ', loss.item())\n",
    "\n",
    "    train_loss = (sum(trn) / len(trn))\n",
    "    train_loss_list.append(train_loss)\n",
    "\n",
    "    net.eval()\n",
    "    val = []\n",
    "    with torch.no_grad() :\n",
    "        for i, (images, labels) in enumerate(data.valid_dl) :\n",
    "            if torch.cuda.is_available():\n",
    "                images = torch.autograd.Variable(images).cuda().float()\n",
    "                labels = torch.autograd.Variable(labels).cuda()\n",
    "            else : \n",
    "                images = torch.autograd.Variable(images).float()\n",
    "                labels = torch.autograd.Variable(labels)\n",
    "\n",
    "            # Forward pass\n",
    "            y_pred = net(images)\n",
    "            \n",
    "            loss = F.cross_entropy(y_pred, labels)\n",
    "            val.append(loss.item())\n",
    "\n",
    "    val_loss = sum(val) / len(val)\n",
    "    val_loss_list.append(val_loss)\n",
    "    val_acc = _get_accuracy(data.valid_dl, net)\n",
    "\n",
    "    print('epoch : ', epoch + 1, ' / ', hyper_params[\"num_epochs\"], ' | TL : ', train_loss, ' | VL : ', val_loss, ' | VA : ', val_acc * 100)\n",
    "\n",
    "    if (val_acc * 100) > min_val :\n",
    "        print('saving model')\n",
    "        min_val = val_acc * 100\n",
    "        torch.save(net.state_dict(), '../saved_models/small_classifier/model1.pt')savename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.cpu()\n",
    "net.load_state_dict(torch.load('../saved_models/small_classifier/model4.pt', map_location = 'cpu'))\n",
    "net.cuda()\n",
    "\n",
    "learn = cnn_learner(data, models.resnet34, metrics = accuracy)\n",
    "learn = learn.load('unfreeze_imagenet_bs64')\n",
    "learn.freeze()\n",
    "\n",
    "print(_get_accuracy(data.valid_dl, net))\n",
    "print(_get_accuracy(data.valid_dl, learn.model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from fastai.vision import *\n",
    "import torch\n",
    "from torchsummary import summary\n",
    "torch.cuda.set_device(0)\n",
    "\n",
    "for repeated in range(0, 1) : \n",
    "    torch.manual_seed(repeated)\n",
    "    torch.cuda.manual_seed(repeated)\n",
    "\n",
    "    # stage should be in 0 to 5 (5 for classifier stage)\n",
    "    hyper_params = {\n",
    "        \"stage\": 5,\n",
    "        \"repeated\": repeated,\n",
    "        \"num_classes\": 10,\n",
    "        \"batch_size\": 64,\n",
    "        \"num_epochs\": 100,\n",
    "        \"learning_rate\": 1e-4\n",
    "    }\n",
    "\n",
    "    path = untar_data(URLs.IMAGENETTE)\n",
    "    tfms = get_transforms(do_flip=False)\n",
    "    data = ImageDataBunch.from_folder(path, train = 'train', valid = 'val', bs = hyper_params[\"batch_size\"], size = 224, ds_tfms = tfms).normalize(imagenet_stats)\n",
    "    \n",
    "    class Flatten(nn.Module) :\n",
    "        def forward(self, input):\n",
    "            return input.view(input.size(0), -1)\n",
    "\n",
    "    def conv2(ni, nf) : \n",
    "        return conv_layer(ni, nf, stride = 2)\n",
    "\n",
    "    class ResBlock(nn.Module):\n",
    "        def __init__(self, nf):\n",
    "            super().__init__()\n",
    "            self.conv1 = conv_layer(nf,nf)\n",
    "\n",
    "        def forward(self, x): \n",
    "            return (x + self.conv1(x))\n",
    "\n",
    "    def conv_and_res(ni, nf): \n",
    "        return nn.Sequential(conv2(ni, nf), ResBlock(nf))\n",
    "\n",
    "    def conv_(nf) : \n",
    "        return nn.Sequential(conv_layer(nf, nf), ResBlock(nf))\n",
    "\n",
    "    net = nn.Sequential(\n",
    "        conv_layer(3, 64, ks = 7, stride = 2, padding = 3),\n",
    "        nn.MaxPool2d(3, 2, padding = 1),\n",
    "        conv_(64),\n",
    "        conv_and_res(64, 128),\n",
    "        conv_and_res(128, 256),\n",
    "        AdaptiveConcatPool2d(),\n",
    "        Flatten(),\n",
    "        nn.Linear(2 * 256, 128),\n",
    "        nn.Linear(128, hyper_params[\"num_classes\"])\n",
    "    )\n",
    "\n",
    "    net.cpu()\n",
    "    filename = '../saved_models/small_stage4/model' + str(repeated) + '.pt'\n",
    "    net.load_state_dict(torch.load(filename, map_location = 'cpu'))\n",
    "\n",
    "    if torch.cuda.is_available() : \n",
    "        net = net.cuda()\n",
    "        print('Model on GPU')\n",
    "\n",
    "    for name, param in net.named_parameters() : \n",
    "        param.requires_grad = False\n",
    "        if name[0] == '7' or name[0] == '8':\n",
    "            param.requires_grad = True\n",
    "        \n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr = hyper_params[\"learning_rate\"])\n",
    "    total_step = len(data.train_ds) // hyper_params[\"batch_size\"]\n",
    "    train_loss_list = list()\n",
    "    val_loss_list = list()\n",
    "    min_val = 0\n",
    "    savename = '../saved_models/small_classifier/model' + str(repeated) + '.pt'\n",
    "    for epoch in range(hyper_params[\"num_epochs\"]):\n",
    "        trn = []\n",
    "        net.train()\n",
    "        for i, (images, labels) in enumerate(data.train_dl) :\n",
    "            if torch.cuda.is_available():\n",
    "                images = torch.autograd.Variable(images).cuda().float()\n",
    "                labels = torch.autograd.Variable(labels).cuda()\n",
    "            else : \n",
    "                images = torch.autograd.Variable(images).float()\n",
    "                labels = torch.autograd.Variable(labels)\n",
    "\n",
    "            y_pred = net(images)\n",
    "\n",
    "            loss = F.cross_entropy(y_pred, labels)\n",
    "            trn.append(loss.item())\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "    #         torch.nn.utils.clip_grad_value_(net.parameters(), 10)\n",
    "            optimizer.step()\n",
    "\n",
    "            if i % 50 == 49 :\n",
    "                print('epoch = ', epoch, ' step = ', i + 1, ' of total steps ', total_step, ' loss = ', loss.item())\n",
    "\n",
    "        train_loss = (sum(trn) / len(trn))\n",
    "        train_loss_list.append(train_loss)\n",
    "\n",
    "        net.eval()\n",
    "        val = []\n",
    "        with torch.no_grad() :\n",
    "            for i, (images, labels) in enumerate(data.valid_dl) :\n",
    "                if torch.cuda.is_available():\n",
    "                    images = torch.autograd.Variable(images).cuda().float()\n",
    "                    labels = torch.autograd.Variable(labels).cuda()\n",
    "                else : \n",
    "                    images = torch.autograd.Variable(images).float()\n",
    "                    labels = torch.autograd.Variable(labels)\n",
    "\n",
    "                # Forward pass\n",
    "                y_pred = net(images)\n",
    "\n",
    "                loss = F.cross_entropy(y_pred, labels)\n",
    "                val.append(loss.item())\n",
    "\n",
    "        val_loss = sum(val) / len(val)\n",
    "        val_loss_list.append(val_loss)\n",
    "        val_acc = _get_accuracy(data.valid_dl, net)\n",
    "\n",
    "        print('epoch : ', epoch + 1, ' / ', hyper_params[\"num_epochs\"], ' | TL : ', train_loss, ' | VL : ', val_loss, ' | VA : ', val_acc * 100)\n",
    "\n",
    "        if (val_acc * 100) > min_val :\n",
    "            print('saving model')\n",
    "            min_val = val_acc * 100\n",
    "            torch.save(net.state_dict(), savename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (ak_fastai)",
   "language": "python",
   "name": "ak_fastai"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
